use polars::prelude::*;
use rust_xlsxwriter::Workbook;
use aes::Aes256;
use cbc::{Encryptor, cipher::{block_padding::Pkcs7, BlockEncryptMut, KeyIvInit}};
use base64::{Engine, engine::general_purpose};

// AES-256-CBC 加密器类型别名
pub type Aes256CbcEnc = Encryptor<Aes256>;

/// 加密单个字符串值
fn encrypt_value(plaintext: &str, key: &[u8; 32], iv: &[u8; 16]) -> Result<String, Box<dyn std::error::Error>> {
    let cipher = Aes256CbcEnc::new(key.into(), iv.into());
    let data = plaintext.as_bytes();
    
    // 创建足够大的缓冲区用于加密和填充
    let mut buf = vec![0u8; data.len() + 16]; // 预留填充空间
    buf[..data.len()].copy_from_slice(data);
    
    let encrypted_data = cipher.encrypt_padded_mut::<Pkcs7>(&mut buf, data.len())
        .map_err(|e| format!("加密失败: {:?}", e))?;
    
    // 转换为base64编码
    Ok(general_purpose::STANDARD.encode(encrypted_data))
}

/// 将DataFrame写入Excel文件
fn write_dataframe_to_excel(df: &DataFrame, output_path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let mut workbook = Workbook::new();
    let worksheet = workbook.add_worksheet();

    // 写入表头
    let headers = df.get_column_names();
    for (col_idx, header) in headers.iter().enumerate() {
        worksheet.write_string(0, col_idx as u16, *header)?;
    }

    // 写入数据
    let height = df.height();
    for row_idx in 0..height {
        for (col_idx, column) in df.get_columns().iter().enumerate() {
            let value = match column.get(row_idx)? {
                AnyValue::Utf8(s) => s.to_string(),
                AnyValue::Binary(b) => String::from_utf8_lossy(b).to_string(),
                other => format!("{}", other),
            };
            worksheet.write_string((row_idx + 1) as u32, col_idx as u16, &value)?;
        }
    }

    workbook.save(output_path)?;
    Ok(())
}

/// 生成包含加密数据的Excel文件
/// 
/// # 参数
/// - output_path: 输出文件路径
/// - encrypt_columns: 需要加密的列名（如["A", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
/// - df: 示例数据DataFrame
pub fn generate_encrypted_excel(
    output_path: &str,
    encrypt_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
    mut df: DataFrame,
) -> Result<(), Box<dyn std::error::Error>> {
    if df.height() == 0 {
        return Err("示例数据不能为空".into());
    }

    // 对需要加密的列进行处理
    for &col_name in encrypt_columns {
        if df.get_column_names().contains(&col_name) {
            // 获取列数据并加密
            let encrypted_values: Result<Vec<String>, Box<dyn std::error::Error>> = df.column(col_name)?
                .utf8()?
                .into_iter()
                .map(|opt_val| {
                    match opt_val {
                        Some(plaintext) => encrypt_value(plaintext, key, iv),
                        None => Ok(String::new()),
                    }
                })
                .collect();
            
            let encrypted_values = encrypted_values?;
            
            // 创建新的Series并替换原列
            let new_series = Series::new(col_name, encrypted_values);
            df.with_column(new_series)?;
        }
    }
    
    // 根据输出路径的扩展名决定输出格式
    if output_path.ends_with(".csv") {
        // 输出为CSV
        let mut file = std::fs::File::create(output_path)?;
        CsvWriter::new(&mut file)
            .include_header(true)
            .finish(&mut df.clone())?;
    } else if output_path.ends_with(".xlsx") {
        // 输出为Excel
        write_dataframe_to_excel(&df, output_path)?;
    } else {
        return Err("Unsupported output format. Use .csv or .xlsx extension.".into());
    }
    
    Ok(())
}

/// 生成默认的示例数据
pub fn create_sample_data() -> Result<DataFrame, Box<dyn std::error::Error>> {
    let col_a = Series::new("A", &["秘密信息1", "秘密信息2", "秘密信息3", "用户密码", "银行卡号"]);
    let col_b = Series::new("B", &["公开信息1", "公开信息2", "公开信息3", "用户名", "持卡人"]);
    let col_c = Series::new("C", &["机密数据1", "机密数据2", "机密数据3", "身份证号", "CVV码"]);
    
    let df = DataFrame::new(vec![col_a, col_b, col_c])?;
    Ok(df)
}

/// 便捷方法：生成包含默认示例数据的加密Excel文件
pub fn generate_sample_encrypted_excel(
    output_path: &str,
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    let sample_data = create_sample_data()?;
    let encrypt_columns = ["A", "C"]; // 加密A列和C列
    
    generate_encrypted_excel(output_path, &encrypt_columns, key, iv, sample_data)?;
    
    println!("已生成加密Excel文件: {}", output_path);
    println!("加密列: {:?}", encrypt_columns);
    println!("文件包含以下数据:");
    println!("- A列(加密): 秘密信息1, 秘密信息2, 秘密信息3, 用户密码, 银行卡号");
    println!("- B列(明文): 公开信息1, 公开信息2, 公开信息3, 用户名, 持卡人");
    println!("- C列(加密): 机密数据1, 机密数据2, 机密数据3, 身份证号, CVV码");
    
    Ok(())
}